import math
from collections import defaultdict
from operator import itemgetter
from util import load_file, save_file, save_to_txt


class ItemCf(object):
    def __init__(self, output_path='./output.txt'):
        self.output_path=output_path

    def __init_data(self, data):
        """
        获得训练所需的数据(user,item)集合
        :param data: 训练集数据
        """
        self.training = dict()
        for user, item, _ in data:
            self.training.setdefault(user, set())
            self.training[user].add(item)

    def train(self, data, matrix_path="./item_matrix.pkl"):
        """
        训练方法,首先加载先用的相似度矩阵进行计算,若不存在则回重新计算相似矩阵
        可以使得预测快一点,在数据不发生变化的前提下
        :param data: 训练数据
        :param matrix_path:矩阵保存路径
        """
        self.data = data
        self.__init_data(data)

        try:
         #   print("load item matrix locally...")
            self.item_like_matrix = load_file(matrix_path)
        except BaseException:
         #   print("load item matrix locally fail ,try to calculate matrix manually")
            self.item_like_matrix = self.get_item_like_matrix()
         #   print("save item matrix ..")
            save_file(matrix_path, self.item_like_matrix)
         #   print("save successfully")

    def get_item_like_matrix(self):
        """
        获得item之间的相似度矩阵
        :return: 计算后的相似度矩阵
        """
        item_like_matrix = dict()
        N = defaultdict(int)

        for u, items in self.training.items():
            for i in items:
                item_like_matrix.setdefault(i, dict())
                N[i] += 1
                for j in items:
                    if i == j:
                        continue
                    item_like_matrix[i].setdefault(j, 0)
                    item_like_matrix[i][j] += 1

        for i, related_items in item_like_matrix.items():
            for j, support in related_items.items():
                item_like_matrix[i][j] = support / math.sqrt(N[i] * N[j])

        return item_like_matrix

    def recommend(self, users, N, K):
        """
        为单个用户推荐的方法
        :param users: 用户的ID
        :param K: 选取K个相似用户
        :param N: 推荐N个电影
        :return: 推进的电影ID集合
        """
        rank = dict()
        ru = self.training[users]
        for r in ru:
            for i, sim in sorted(self.item_like_matrix.get(r, {}).items()
                    , key=itemgetter(1), reverse=True)[:K]:
                if i in ru:
                    continue
                rank.setdefault(i, 0.)
                rank[i] += sim
        return dict(sorted(rank.items(), key=itemgetter(1), reverse=True)[:N])

    def recommend_test(self, users, N, K):
        """
        在测试集合上推荐,调用recommend方法
        :param users: 用户的ID集合
        :param K: 选取K个相似用户
        :param N: 推荐N个电影
        :return: {(user_1,{ID_set}),(user_2,{ID_set_2})..}
        """

        rec_set = dict()
        for user in users:
            user_rec = list(self.recommend(user, N, K).keys())
            rec_set[user] = user_rec
        save_to_txt(data = rec_set,path=self.output_path)
        return rec_set
